--[[
Indentation checker for Lua files, written in Lua

Copyright 2022 Pedro Gimeno Fortea

Permission to use, copy, modify, and/or distribute this software for any
purpose with or without fee is hereby granted, provided that the above
copyright notice and this permission notice appear in all copies.

THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
PERFORMANCE OF THIS SOFTWARE. 
--]]

local pos = 1
local state = 0
local token = {}
local prog
local progLen
local lineNo
local a, b
local tabSize = 8
local nextIndent = 0
local indentLevel

-- Aux functions

local function showWarning(msg)
  io.stderr:write("Warning at line ", tostring(lineNo), ": ", msg, "\n")
  io.stderr:flush()
end

local function matchesPattern(pat)
  a, b = prog:find("^" .. pat, pos)
  return a
end

local function findIndent(str)
  local indent = 0
  for i = 1, #str do
    if str:sub(i, i) == "\t" then
      indent = (indent + tabSize) - (indent + tabSize) % tabSize
    else
      indent = indent + 1
    end
  end
  return indent
end

local function makeHash(t)
  local r = {}
  for i = 1, #t do r[t[i]] = true end
  return r
end

-- Table for conversion of backslash escapes
local xlat = {a="\a", b="\b", f="\f", n="\n", r="\r", t="\t", v="\v",
  ["'"]="'", ["\\"]="\\", ['"']='"', ["\n"]="\n"}

-- Table of reserved words (as a hash)
local reserved = makeHash{ "and", "break", "do", "else", "elseif", "end",
  "false", "for", "function", "if", "in", "local", "nil", "not", "or",
  "repeat", "return", "then", "true", "until", "while" }

-- Convert a string to its actual value

local function parseString()
  local term = prog:sub(pos, pos)
  local s = ""
  local spanMatch = "^[^\n\\"..term.."]+"
  repeat
    pos = pos + 1
    local span = prog:match(spanMatch, pos)
    if span then
      s = s .. span
      pos = pos + #span
    end
    if pos > progLen then
      showWarning("Unterminated string")
      return s
    end
    local c = prog:sub(pos, pos)
    if c == term then
      pos = pos + 1
      return s
    end
    if c == '\n' then
      showWarning("Unterminated string")
      return s
    end
    if c == '\\' then
      pos = pos + 1
      if pos > progLen then
        showWarning("Unterminated string")
        return s
      end
      c = prog:sub(pos, pos)
      if not c:match("^[\"'\\0-9abfnrtvx\n]$") then
        showWarning("Unknown escape in string, treating as single char")
      end
      if c:find("[0-9]") then
        local v = prog:match("[0-9][0-9]?[0-9]?", pos)
        v = tonumber(v)
        if v > 255 then
          showWarning("Invalid decimal code in string");
        else
          s = s .. string.char(v)
          pos = pos + #v - 1
        end
      elseif c == 'x' then
        -- \X is not valid, so don't check uppercase
        local v = prog:match("x[0-9A-Fa-f][0-9A-Fa-f]?", pos)
        if not v then
          showWarning("Invalid hex code in string");
        else
          s = s .. string.char(tonumber("0" .. v))
          pos = pos + #v - 1
        end
      else
        s = s .. xlat[c]
      end
    else
      s = s .. c
    end
  until false
end

-- Scanner

local function getToken()
  while pos <= progLen do
    if matchesPattern("%-%-%[=*%[") then
      -- comment that must be matched by ]==...==]
      local start = b + 1
      local len = b - a - 3 -- -4 + 1 because it's inclusive
      local term = "]" .. ("="):rep(len) .. "]"
      a, b = prog:find(term, b + 1, true)
      local str
      if not a then
        showWarning("Unterminated comment")
        pos = progLen + 1
        str = prog:sub(start, pos - 1)
      else
        pos = b + 1
        str = prog:sub(start, pos - len - 3)
      end
      for i in str:gmatch("\n") do
        lineNo = lineNo + 1
      end
    elseif matchesPattern("%-%-") then
      -- single-line comment
      a, b = prog:find("\n", pos + 2)
      if a then
        pos = b -- to allow the \n to be detected and accounted for
      else
        pos = progLen + 1
      end
    elseif matchesPattern("[\r\t\f ]+") then
      pos = b + 1

    elseif matchesPattern("\n[\t ]*") then
      lineNo = lineNo + 1
      pos = b + 1
      -- Find indent level for next token
      nextIndent = findIndent(prog:sub(a + 1, b))

    else
      -- Not whitespace or comments - we're returning a token
      indentLevel = nextIndent

      if matchesPattern("%[=*%[") then
        -- string that must be matched by ]==...==]
        local start = b + 1
        local len = b - a - 1 -- -2 + 1 because it's inclusive
        local term = "]" .. ("="):rep(len) .. "]"
        a, b = prog:find(term, start, true)
        token[1] = "STRING"
        if not a then
          showWarning("Unterminated string")
          pos = progLen + 1
          token[2] = prog:sub(start, pos - 1)
        else
          pos = b + 1
          token[2] = prog:sub(start, pos - len - 3)
        end
        for i in token[2]:gmatch("\n") do
          lineNo = lineNo + 1
        end
        return
      elseif matchesPattern("[\"']") then
        token[1] = "STRING"
        token[2] = parseString()
        return
      elseif matchesPattern("0[Xx][0-9a-fA-F]+[Pp][-+]?[0-9]+")
          or matchesPattern("0[Xx][0-9a-fA-F]+")
          or matchesPattern("[0-9]+%.[0-9]*[Ee][-+]?[0-9]+")
          or matchesPattern("%.[0-9]+[Ee][-+]?[0-9]+")
          or matchesPattern("[0-9]+%.[0-9]*")
          or matchesPattern("%.[0-9]+")
          or matchesPattern("[0-9]+[Uu]?[Ll][Ll]")
          or matchesPattern("[0-9]+")
      then
        pos = b + 1
        token[1] = "NUM"
        token[2] = prog:sub(a, b)
        return
      elseif matchesPattern("[A-Za-z_][A-Za-z0-9_]*") then
        pos = b + 1
        local ident = prog:sub(a, b)
        if reserved[ident] then
          token[1] = ident
          token[2] = nil
        else
          token[1] = "IDENT"
          token[2] = prog:sub(a, b)
        end
        return
      else -- symbol
        token[2] = nil
        -- multi-symbol tokens
        if matchesPattern("==") or matchesPattern(">=") or matchesPattern("<=")
          or matchesPattern("~=") or matchesPattern("%.%.%.")
          or matchesPattern("%.%.")
          or matchesPattern(".") -- catch-all
        then
          pos = b + 1
          token[1] = prog:sub(a, b)
        end
        return
      end
    end
  end
  token[1] = "EOF"
  token[2] = nil
end

local function printToken()
  print(string.format("{ %s, %s } %d", token[1], 
    token[1] == "STRING" and ("%q"):format(tostring(token[2])):gsub("\n","n")
    or token[2], indentLevel))
end

local function dumpTokens()
  repeat
    printToken()
    if pos > progLen then
      break
    end
    getToken()
  until false
end

local function skipToToken(tok)
  if tok ~= "EOF" then
    repeat
      getToken()
    until token[1] == tok or token[1] == "EOF"
    if token[1] ~= tok then
      showWarning("Unexpected EOF")
    end
  end
end

local blockTokens = makeHash{ "do", "if", "while", "for", 
  "repeat", "function" }
local termTokens = makeHash{ "end", "elseif", "else", "until" }

local function parseBlock(term)
  repeat
    local tok = token[1]
    if tok == "EOF" then
      if term ~= "EOF" then
        showWarning("Unexpected EOF looking for block terminator")
      end
      return
    end

    if term == "if" then
      if tok == "elseif" or tok == "end" or tok == "else" then
        return
      end
      if termTokens[tok] then
        showWarning("Unbalanced token: " .. tok)
        return
      end
    else
      if tok == term then
        return
      end
      if termTokens[tok] then
        showWarning("Unbalanced token: " .. tok)
        return
      end
    end
    if blockTokens[tok] then
      local indentAtBlockStart = indentLevel -- in case it's a "do"
      getToken()
      if tok == "for" or tok == "while" then
        skipToToken("do")
        if token[1] == "do" then -- may be EOF too
          getToken()
          parseBlock("end")
          if indentLevel ~= indentAtBlockStart then
            showWarning("Misindented end ("..tok..")")
          end
        end
      elseif tok == "repeat" then
        parseBlock("until")
        if indentLevel ~= indentAtBlockStart then
          showWarning("Misindented until (repeat)")
        end
      elseif tok == "function" then
        skipToToken(")")
        if token[1] == ")" then
          getToken()
          parseBlock("end")
          if indentLevel ~= indentAtBlockStart then
            showWarning("Misindented end (function)")
          end
        end
      elseif tok == "if" then
        repeat
          skipToToken("then")
          if token[1] == "then" then
            getToken()
            parseBlock("if") -- elseif,else,end
            if indentLevel ~= indentAtBlockStart then
              showWarning("Misindented "..token[1].." (if)")
            end
          end
        until token[1] ~= "elseif"
        if token[1] == "else" then
          indentAtBlockStart = indentLevel
          getToken()
          parseBlock("end")
          if indentLevel ~= indentAtBlockStart then
            showWarning("Misindented end (else)")
          end
        end
      else -- "do"
        assert(tok == "do")
        parseBlock("end")
        if indentLevel ~= indentAtBlockStart then
          showWarning("Misindented end (do)")
        end
      end
      getToken()
    else
      getToken()
    end
  until false
end

local function findMissingEnd(prog)
  progLen = #prog
  lineNo = 1
  pos = 1
  if matchesPattern("[ \t]*") then -- always true, possibly empty
    nextIndent = findIndent(prog:sub(a, b))
  end
  getToken()
  --dumpTokens()
  parseBlock("EOF")
end

do
  local argcopy = {}
  for i, v in ipairs(arg) do argcopy[i] = v end
  
  local i = 1
  local showHelp = false
  while argcopy[i] ~= nil do
    if argcopy[i] == "-t" then
      table.remove(argcopy, i)
      if argcopy[i] ~= nil then
        tabSize = tonumber(argcopy[i])
        if tabSize == nil then
          tabSize = 8
        end
        table.remove(argcopy, i)
      end
    elseif argcopy[i] == "-h" or argcopy[i] == "--help" then
      showHelp = true
    end
    i = i + 1
  end
  
  if showHelp then
    print("Usage: lua "..arg[0].." [-t tabsize] file1 [file2 ...]")
    return
  end
  
  for i = 1, #argcopy do
    io.stderr:write("**** ", argcopy[i], "\n")
    local f = io.open(argcopy[i], "rb")

    -- Little hack to deal with CRLF: preprocess the input file to convert all 
    -- CRLFs to LFs. This avoids problems like a backslash line continuation
    -- inside a string in systems with CRLF line terminators to escape the CR
    -- only and err on the LF for not recognizing CRLF as a single line 
    -- terminator. An actual patch for this behaviour would be 
    -- disproportionately more complex.
    prog = f:read("*a"):gsub("\r\n", "\n")
    f:close()
  end
end

findMissingEnd(prog)
